package com.rui.study.commons.support;

import com.rui.study.commons.model.BizException;
import com.rui.study.commons.model.StatusCode;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @program: study
 * @description:
 * @author: Yaowr
 * @create: 2018-11-20 14:31
 **/
@Service
@RefreshScope
public class ZKLock {

    private static final String LOCK_PATH = "/lock/";

    private ConcurrentHashMap<Thread, Map<String,InterProcessMutex>> threadMutexMap = new ConcurrentHashMap<>();

    @Value("${zk.path}")
    private String zkPath;

    private CuratorFramework client;

    @PostConstruct
    private void init() {
        client = CuratorFrameworkFactory.newClient(zkPath,
                new ExponentialBackoffRetry(1000, 3));
        client.start();
    }

    public void lock(String method) {
        try {
            fetchMutex(method).acquire();
        } catch (Exception e) {
            throw new BizException(StatusCode.SYSTEM_EXCEPTION, e.getMessage());
        }
    }

    public void unlock(String method) {
        try {
            fetchMutex(method).release();
            clear(method);
        } catch (Exception e) {
            throw new BizException(StatusCode.SYSTEM_EXCEPTION, e.getMessage());
        }
    }

    private void clear(String method) {
        /* TODO 清理缓存的mutex */
    }


    private InterProcessMutex fetchMutex(String method) {
        Thread current = Thread.currentThread();
        Map<String, InterProcessMutex> mutexMap = threadMutexMap.get(current);
        if (mutexMap == null) {
            mutexMap = new HashMap<>();
            threadMutexMap.put(current, mutexMap);
        }

        InterProcessMutex mutex = mutexMap.get(method);
        if (mutex == null) {
            mutex = new InterProcessMutex(client, LOCK_PATH + method);
            mutexMap.put(method, mutex);
        }
        return mutex;
    }
}
